# 更新 apt /安装基本工具
apt-get update
apt-get install -y git cmake build-essential ninja-build libssl-dev libffi-dev \
    python3-dev python3-pip libopenblas-dev wget

# 安装 Python 3.10 和 dev 头文件
apt-get install -y python3.10 python3.10-dev

# 检查 Python 3.10 是否可用
python3.10 --version || { echo "python3.10 not installed"; exit 1; }

# 安装 pip for Python 3.10（如果还没安装）
wget https://bootstrap.pypa.io/get-pip.py -O get-pip.py
python3.10 get-pip.py

# 检查 pip
python3.10 -m pip --version || { echo "pip for python3.10 failed"; exit 1; }

# 卸载旧版本包（如果存在于 Python3.10 环境中）
python3.10 -m pip uninstall -y torch torchvision torchaudio nemo_toolkit nemo huggingface_hub

# 克隆 PyTorch 源码（如果还没克隆）
cd /root
if [ ! -d pytorch ]; then
  git clone --recursive https://github.com/pytorch/pytorch.git
fi
cd pytorch

# 可选：切换分支（如果你知道适合版本）
# git checkout v2.7.0
git submodule sync
git submodule update --init --recursive

# 设置环境变量以支持 sm_120 架构
export TORCH_CUDA_ARCH_LIST="120;100;90;86;80;75;70;60"
# 让 CMake 找到正确的 Python 前缀
export CMAKE_PREFIX_PATH="$(python3.10 -c 'import sys; print(sys.prefix)')"

# 编译 & 安装 PyTorch — 用 python3.10
python3.10 setup.py install

# 验证安装结果
python3.10 - << 'EOF'
import torch
print("Torch version:", torch.__version__)
print("CUDA version:", torch.version.cuda)
print("CUDA available:", torch.cuda.is_available())
print("Supported arch list:", torch.cuda.get_arch_list())
EOF
